"""
Provide basic traversal operations.
"""

from gremlin_python.process.graph_traversal import GraphTraversal
from gremlin_python.process.traversal import TextP
from gremlin_python.process.graph_traversal import __


def q_node_with_code(g: GraphTraversal, code: str) -> GraphTraversal:
    """
    Find all nodes with the given code, exact match.
    """
    return g.has("code", code)


def q_node_contains_code(
    g: GraphTraversal, code: str, position="contains"
) -> GraphTraversal:
    """
    Find all nodes with the given code, with contains, start, or end.
    """
    if position == "contains":
        return g.has("code", TextP.containing(code))
    elif position == "start":
        return g.has("code", TextP.starting_with(code))
    elif position == "end":
        return g.has("code", TextP.ending_with(code))
    else:
        raise ValueError("Invalid position")


def q_dfg_paths(g: GraphTraversal) -> GraphTraversal:
    """
    Find all DFG paths starting from the given nodes.
    """
    return g.repeat(__.out("dfg").simplePath()).until(__.not_(__.out("dfg"))).path()
